#include "opencv2/opencv.hpp"
#include <QDebug>
#include <QString>
#include "aruco/aruco.h"
#include "xcorner/XCornerFinder.h"
#include "colormarker.h"
#include <QFile>
#include "point_fit/pointfit.h"

/*
    This project file is to demonstrate the calibrated stereo camera system
    You can use this project to
    1. Detect ID marker with the built-in ArUco toolkit
    2. Detect colour cross markers

    Components correspondence for both of them will be calculated; also the
    estimated 3D position is calculated and printed on the left screen.



*/

using namespace std;
using namespace cv;
using namespace aruco;

#define SHOW_DISP        false // run stereo matching function stereoVAR and display the disparity map
#define DETECT_ID_MARKER false // detect ID markers by the ArUco library
#define DETECT_CX_MARKER true  // detect colour markers
#define UNIT_SIZE        23.6f // the calibration chessboard square size in Millimetres
#define MARKER_WINDOW    15    // the searching window for the colour marker
#define MAX_X_DISPARITY  200   // maximum horizontal disparity for determine CM correspondence
#define MIN_X_DISPARITY  5
#define MAX_Y_DISPARITY  4

// number of desired correspondent colour markers to be detected
// if the number of markers equals to this number, the capturing will be freezed and ask for
// whether or not to save the marker positions and the current image frame
#define MARKERS_TO_DETECT 5
#define FRAME_TO_FREEZE   20

int main( int, char** )
{
    VideoCapture cap0, cap1;
    if( cap0.open( 0 ) && cap1.open( 1 ) )
    {
        qDebug() << "Cameras #0 & #1 detected and opened.";
    }
    else
    {
        qDebug() << "Please ensure two cameras are plugged in.";
    }

    Mat img0, img1, rimg0, rimg1, crimg0, crimg1;
    Mat grimg0, grimg1, grimg2, grimg3;
    Mat rblend, blend;
    Mat disp, disp8;
    Mat mask;

    Rect valid( 40, 45, 540, 385 );

    cap0 >> img0;
    cap1 >> img1;

    Size imageSize = Size( 640, 480 );

    /* Read in pre-calibrated stereo camera configurations */
    FileStorage fs_in( "intrinsics.yml", FileStorage::READ );
    FileStorage fs_ex( "extrinsics.yml", FileStorage::READ );

    Mat cameraMatrix[2], distCoeffs[2], R[2], P[2], Q;
    Mat rmap[2][2];
    //Rect validRoi[2];
    fs_in["M1"] >> cameraMatrix[0];
    fs_in["M2"] >> cameraMatrix[1];
    fs_in["D1"] >> distCoeffs[0];
    fs_in["D2"] >> distCoeffs[1];

    fs_ex["R1"] >> R[0];
    fs_ex["R2"] >> R[1];
    fs_ex["P1"] >> P[0];
    fs_ex["P2"] >> P[1];
    fs_ex["Q"]  >> Q;

    initUndistortRectifyMap(cameraMatrix[0], distCoeffs[0], R[0], P[0], imageSize, CV_16SC2, rmap[0][0], rmap[0][1]);
    initUndistortRectifyMap(cameraMatrix[1], distCoeffs[1], R[1], P[1], imageSize, CV_16SC2, rmap[1][0], rmap[1][1]);

    /* Recitify images */
    remap( img0, rimg0, rmap[0][0], rmap[0][1], CV_INTER_LINEAR );
    remap( img1, rimg1, rmap[1][0], rmap[1][1], CV_INTER_LINEAR );

    bitwise_and( rimg0, rimg1, mask );

    CameraParameters camparam[2];
    MarkerDetector mdetector;
    vector<Marker> markers0, markers1;

    Marker mymarker;
    int markersize = 200, markerid = 50;
    Mat markerimg;

    if( DETECT_ID_MARKER )
    {
        mdetector.setThresholdMethod( MarkerDetector::ADPT_THRES );
        markerimg = mymarker.createMarkerImage( markerid, markersize );
        namedWindow( "marker", 1 );
        imshow( "marker", markerimg );
    }

    camparam[0].readFromXMLFile( "cam_L.yml" );
    camparam[0].resize( img0.size() );
    camparam[1].readFromXMLFile( "cam_R.yml" );
    camparam[1].resize( img0.size() );

    // block matching
    StereoSGBM sgbm;
    sgbm.preFilterCap = 63;
    sgbm.SADWindowSize = 3;
    sgbm.P1 = 8*3*sgbm.SADWindowSize*sgbm.SADWindowSize;
    sgbm.P2 = 32*3*sgbm.SADWindowSize*sgbm.SADWindowSize;
    sgbm.minDisparity = 0;
    sgbm.numberOfDisparities = ((imageSize.width/8) + 15) & -16;
    sgbm.uniquenessRatio = 10;
    sgbm.speckleWindowSize = 100;
    sgbm.speckleRange = 32;
    sgbm.disp12MaxDiff = 1;
    sgbm.uniquenessRatio = 10;

    StereoVar var;
    var.levels = 6;
    var.pyrScale = 0.6;
    var.nIt = 3;
    var.minDisp = -sgbm.numberOfDisparities;
    var.maxDisp = 0;
    var.poly_n = 3;
    var.poly_sigma = 0.0;
    var.fi = 5.0f;
    var.lambda = 0.1;
    var.penalization = var.PENALIZATION_TICHONOV;
    var.cycle = var.CYCLE_V;
    var.flags = var.USE_SMART_ID | var.USE_INITIAL_DISPARITY | 1 * var.USE_MEDIAN_FILTERING ;

    ColorMarkerDetector cmdetector( 640, 480 );

    int frame_counter = 1;
    double err_sum = 0;
    QFile fout( "coordinates.txt" );
    if( !fout.open( QIODevice::ReadWrite | QIODevice::Truncate ) )
    {
        qDebug() << "can not output file";
    }
    QTextStream outStream( &fout );
    int     to_show_counter = 0;
    string  info_to_show;

    vector<Point3f> originView, thisRes;
    for( int frame_saved = 0;;)
    {
        thisRes.clear();
        cap0 >> img0;
        cap1 >> img1;

        if( DETECT_ID_MARKER )
        {
            mdetector.detect( img0, markers0, camparam[0], markersize );
            mdetector.detect( img1, markers1, camparam[1], markersize );

            // we draw only the first marker
            if( markers0.size() && markers1.size()  )
            {
                markers0.at(0).draw( img0, Scalar( 255, 0, 0 ) );
                markers1.at(0).draw( img1, Scalar( 255, 0, 0 ) );

                CvDrawingUtils::draw3dAxis( img0, markers0.at( 0 ), camparam[0] );
                // rectify the points so that they are correctly aligned
                // similar to rectify map
                undistortPoints( markers0[0], markers0[0], cameraMatrix[0], distCoeffs[0], R[0], P[0] );
                undistortPoints( markers1[0], markers1[0], cameraMatrix[1], distCoeffs[1], R[1], P[1] );

                Point pl_0 = markers0.at(0)[0], pr_0 = (Point)markers1.at(0)[0];
                Point pl_1 = markers0.at(0)[1], pr_1 = (Point)markers1.at(0)[1];
                Point pl_2 = markers0.at(0)[2], pr_2 = (Point)markers1.at(0)[2];
                Point pl_3 = markers0.at(0)[3], pr_3 = (Point)markers1.at(0)[3];

                Mat pnts( 4, 1, CV_32FC3, Scalar::all( 0 ) );
                pnts.row( 0 ) = Scalar( pl_0.x, pl_0.y, abs( pl_0.x - pr_0.x ) ) ;
                pnts.row( 1 ) = Scalar( pl_1.x, pl_1.y, abs( pl_1.x - pr_1.x ) ) ;
                pnts.row( 2 ) = Scalar( pl_2.x, pl_2.y, abs( pl_2.x - pr_2.x ) ) ;
                pnts.row( 3 ) = Scalar( pl_3.x, pl_3.y, abs( pl_3.x - pr_3.x ) ) ;
                perspectiveTransform( pnts, pnts, Q );
                pnts = pnts * UNIT_SIZE;
                QString sp[4], info[4];
                double diff_sum = 0;
                for( int i = 0; i < 4; i ++ )
                {
                    Vec3f t_row = pnts.at<Vec3f>( i, 0 );
                    sp[i] = "(" + QString::number( t_row[0], 'f', 1 );
                    sp[i] += ", " + QString::number( t_row[1], 'f', 1 );
                    sp[i] += ", " + QString::number( t_row[2], 'f', 1 ) + ")";
                    putText( img0, sp[i].toStdString(), markers0.at(0)[i],
                             CV_FONT_HERSHEY_TRIPLEX, 0.4, Scalar( 0, 0, 255 ), 2 );

                    Mat diff = pnts.row( i ) - pnts.row( ( i + 1 ) % 4 );
                    diff_sum += sqrt( diff.dot( diff ) );
                    info[i] = "P" + QString::number( i ) + " to P" + QString::number( ( i + 1 ) % 4 )
                            + " = " + QString::number( sqrt( diff.dot( diff ) ) );
                    putText( img0, info[i].toStdString(), Point( 20, 20 + 20 * i ), CV_FONT_HERSHEY_TRIPLEX, 0.5,
                             Scalar( 0, 0, 0 ), 2 );
                    putText( img0, info[i].toStdString(), Point( 20, 20 + 20 * i ), CV_FONT_HERSHEY_TRIPLEX, 0.5,
                             Scalar( 255, 255, 255 ), 1 );
                }
                QString info_sum = "Avg. (cm) = " + QString::number( diff_sum / 4 );
                err_sum += diff_sum / 4 - 5.2;
                        + " for " + QString::number( frame_counter ) + " frames.";
                putText( img0, info_sum.toStdString(), Point( 20, 20 + 20 * 4 ), CV_FONT_HERSHEY_TRIPLEX, 0.5,
                         Scalar( 0, 0, 0 ), 2 );
                putText( img0, info_sum.toStdString(), Point( 20, 20 + 20 * 4 ), CV_FONT_HERSHEY_TRIPLEX, 0.5,
                         Scalar( 255, 0, 255 ), 1 );
                frame_counter ++;
            }
        }

        int indeces [1024][2], counter = 0;
        vector<ColorMarker> colorMarkers0;
        vector<ColorMarker> colorMarkers1;
        vector<Point2f> cMarkers0, cMarkers1;
        Mat c_pnts;
        QString outstring;
        QString filename;
        if ( DETECT_CX_MARKER )
        {
            cmdetector.detect( img0, colorMarkers0, MARKER_WINDOW );
            cmdetector.detect( img1, colorMarkers1, MARKER_WINDOW );
            if( colorMarkers0.size() && colorMarkers1.size() )
            {
                cMarkers0.clear();
                cMarkers1.clear();
                for( int i = 0; i < colorMarkers0.size(); i ++ )
                {
                    cMarkers0.push_back( colorMarkers0.at( i ) );
                }
                for( int i = 0; i < colorMarkers1.size(); i ++ )
                {
                    cMarkers1.push_back( colorMarkers1.at( i ) );
                }

                undistortPoints( cMarkers0, cMarkers0, cameraMatrix[0], distCoeffs[0], R[0], P[0] );
                undistortPoints( cMarkers1, cMarkers1, cameraMatrix[1], distCoeffs[1], R[1], P[1] );

                c_pnts = Mat( colorMarkers0.size(), 1, CV_32FC3, Scalar::all( 0 ) );

                for( int i = 0; i < colorMarkers0.size(); i ++ )
                {
                    ColorMarker lmarker = colorMarkers0.at( i );
                    Point2f lpnt = cMarkers0.at( i );
                    double min = 1000000;
                    int min_id = -1;
                    for( int j = 0; j < colorMarkers1.size(); j ++ )
                    {
                        ColorMarker rmarker = colorMarkers1.at( j );
                        Point2f rpnt = cMarkers1.at( j );
                        if(     lpnt.x < rpnt.x
                                || abs( lpnt.y - rpnt.y ) > MAX_Y_DISPARITY
                                || abs( lpnt.x - rpnt.x ) < MIN_X_DISPARITY
                                || abs( lpnt.x - rpnt.x ) > MAX_X_DISPARITY )
                            continue;
                        double hist_diff = compareHist( lmarker.hist, rmarker.hist, CV_COMP_BHATTACHARYYA );
                        if( hist_diff < min )
                        {
                            min    = hist_diff;
                            min_id = j;
                        }
                    }
                    if ( min_id != -1 )
                    {
                        Point2f rpnt = cMarkers1.at( min_id );
                        c_pnts.row( counter ) = Scalar( lpnt.x, lpnt.y,
                                                       abs( lpnt.x - rpnt.x ) ) ;
                        indeces[counter][0] = i;
                        indeces[counter][1] = min_id;
                        counter ++;
                    }
                }
                perspectiveTransform( c_pnts, c_pnts, Q );
                c_pnts = c_pnts * UNIT_SIZE;
                QString cm_qstr;
                filename = QString::number( frame_saved ) + ".png";
                outstring = "framefile: " + filename + "\n";
                for( int i = 0; i < counter; i ++ )
                {
                    Vec3f t_row = c_pnts.at<Vec3f>( i, 0 );
                    thisRes.push_back( Point3f( t_row[0], t_row[1], t_row[2] ) );
                    cm_qstr = QString::number( i ) + "";
                    cm_qstr += "(" + QString::number( t_row[0], 'f', 1 );
                    cm_qstr += ", " + QString::number( t_row[1], 'f', 1 );
                    cm_qstr += ", " + QString::number( t_row[2], 'f', 1 ) + ")";
                    putText( img0, cm_qstr.toStdString(), colorMarkers0.at( indeces[i][0] ),
                             CV_FONT_HERSHEY_TRIPLEX, 0.7, Scalar( 0, 0, 0 ), 2, CV_AA );
                    putText( img0, cm_qstr.toStdString(), colorMarkers0.at( indeces[i][0] ),
                             CV_FONT_HERSHEY_TRIPLEX, 0.7, Scalar( 0, 0, 255 ), 1, CV_AA );
                    outstring += "    " + QString::number( i ) + " ("
                            + QString::number( t_row[0] ) + ","
                            + QString::number( t_row[1] ) + ","
                            + QString::number( t_row[2] ) + ") sig("
                            + QString::number( colorMarkers0.at( indeces[i][0] ).hue ) + ")\n";
                }
                //qDebug() << outstring << counter;
                Mat R, T;
                if( counter == MARKERS_TO_DETECT && frame_saved )
                {
                    PointFitter::fit( thisRes, originView, R, T );
                    cout << "R: \n" << R << "\nT: \n" << T << endl;
                }
            }
        }
        /* Recitify images */
        remap( img0, rimg0, rmap[0][0], rmap[0][1], CV_INTER_LINEAR );
        remap( img1, rimg1, rmap[1][0], rmap[1][1], CV_INTER_LINEAR );

        addWeighted( rimg0, 0.5, rimg1, 0.5, 0, rblend );
        // draw color marker correspondences
        if( DETECT_CX_MARKER )
        {
            for( int i = 0; i < counter; i ++ )
            {
                Point cmarker0 = cMarkers0.at( indeces[i][0] );
                Point cmarker1 = cMarkers1.at( indeces[i][1] );
                line( rblend, cmarker0, cmarker1, Scalar( 255, 0, 0 ), 2 );
            }
        }
        // draw disparity map use stereoVAR method
        if( SHOW_DISP )
        {
            cvtColor( rimg0, grimg0, CV_BGRA2GRAY );
            cvtColor( rimg1, grimg1, CV_BGRA2GRAY );

            grimg0.convertTo( grimg2, 32 );
            grimg1.convertTo( grimg3, 32 );
            var( rimg0, rimg1, disp );
            disp.convertTo( disp8, CV_8U );

            imshow( "disparity", disp8 );
        }
        crimg0 = rimg0( valid );
        crimg1 = rimg1( valid );

        addWeighted( crimg0, 0.5, crimg1, 0.5, 0, blend );

        Mat pair( 480, 640 * 2 + 1, CV_8UC3, Scalar::all( 0 ) );
        Mat rpair( 385, 540 * 2 + 1, CV_8UC3, Scalar::all( 0 ) );
        Mat pair_L  = pair( Rect( 0, 0, 640, 480 ) );
        Mat pair_R  = pair( Rect( 641, 0, 640, 480 ) );
        Mat rpair_L = rpair( Rect( 0, 0, 540, 385 ) );
        Mat rpair_R = rpair( Rect( 541, 0, 540, 385 ) );

        add( pair_L, img0, pair_L );
        add( pair_R, img1, pair_R );

        add( rpair_L, crimg0, rpair_L );
        add( rpair_R, crimg1, rpair_R );

        //imshow( "webcam views unrectified", pair );
        //imshow( "img0", img0 );
        //imshow( "img1", img1 );
        //imshow( "rectified blend", rblend );
        imshow( "rectified valid blend", blend );

        if( DETECT_CX_MARKER )
        {
            for( int i = 0; i < counter; i ++ )
            {
                Point cmarker0 = cMarkers0.at( indeces[i][0] );
                Point cmarker1 = cMarkers1.at( indeces[i][1] );
                cmarker0.x -= 40; cmarker0.y -= 45;
                cmarker1.x -= 40; cmarker1.y -= 45;
                cmarker1.x += 541;
                line( rpair, cmarker0, cmarker1, Scalar( 0, 0, 0 ), 2, CV_AA );
                line( rpair, cmarker0, cmarker1, Scalar( 255, 0, 0 ), 1, CV_AA );
            }
        }

        if( to_show_counter )
        {

            to_show_counter --;
            if( to_show_counter > 5 )
                putText( rpair, info_to_show,  Point( 10, 60 ),
                         FONT_HERSHEY_COMPLEX, 0.6, Scalar( 0, 255, 255 ), 1, CV_AA );
        }
        else
        {
            to_show_counter = 0;
        }

        if( counter == MARKERS_TO_DETECT && to_show_counter == 0  )
        {
            putText( rpair, "Save these marker positions?",  Point( 10, 40 ),
                     CV_FONT_HERSHEY_TRIPLEX , 0.6, Scalar( 0, 0, 0 ), 2, CV_AA );
            putText( rpair, "Save these marker positions?",  Point( 10, 40 ),
                     CV_FONT_HERSHEY_TRIPLEX , 0.6, Scalar( 0, 0, 255 ), 1, CV_AA );
        }

        putText( rpair, "Left",  Point( 10, 20 ),
                 CV_FONT_HERSHEY_TRIPLEX , 0.4, Scalar( 0, 255, 255 ), 1, CV_AA );
        putText( rpair, "Right", Point( 10 + 541, 20 ),
                 CV_FONT_HERSHEY_TRIPLEX , 0.4, Scalar( 0, 255, 255 ), 1, CV_AA );

        imshow( "rectified & cropped pair", rpair );
        if( waitKey( 30 ) == 'q' && counter != MARKERS_TO_DETECT )
            break;

        if( counter == MARKERS_TO_DETECT && to_show_counter == 0 )
        {
            if( waitKey( 0 ) == 'y' )
            {
                info_to_show    = "Coordinates saved";
                outStream << outstring << "\n";
                imwrite( filename.toStdString(), rpair  );
                frame_saved ++;
                originView.clear();
                for( int i = 0; i < MARKERS_TO_DETECT; i ++ )
                {
                    originView.push_back( thisRes.at( i ) );
                }
            }
            else
            {
                info_to_show    = "Abandoned";
            }
            to_show_counter = FRAME_TO_FREEZE;
        }
    }
    fout.flush();
    fs_in.release();
    fs_ex.release();
    return 0;
}
